package client;

import java.io.IOException;
import java.math.BigInteger;
import java.net.Socket;
import java.net.UnknownHostException;
import java.security.AlgorithmParameterGenerator;
import java.security.AlgorithmParameters;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.util.Scanner;

import javax.crypto.KeyAgreement;
import javax.crypto.SecretKey;
import javax.crypto.spec.DHParameterSpec;

import messages.DHMessage;
import messages.DHResponseMessage;
import messages.LoginMessage;
import messages.LoginResponseMessage;
import messages.Message;
import util.DHGenerator;
import util.Network;


public class ClientMain {

	/**
	 * @param args
	 */
	private static Socket sock;
	private static final String salt = "38&!";
	
	public static void main(String[] args) {
		try {
			String username, password;

			if(args.length == 1 && args[0].equals("-test")) {
				sock = new Socket("localhost", 25011);
				username = "user1";
				password = salt + "password1";
			} else {
				Scanner scan = new Scanner(System.in);
				System.out.print("Enter the address of the server:  ");
				String address = scan.nextLine();
				
				System.out.print("Enter the port number:  ");
				String port = scan.nextLine();
				sock = new Socket(address, Integer.parseInt(port));
				
				System.out.print("Enter your username:  ");
				username = scan.nextLine();
				System.out.print("Enter your password:  ");
				password = scan.nextLine();
				password = salt + password;
			}
			
			byte[] passBytes = password.getBytes("UTF-8");
			
			ClientReadThread crt = new ClientReadThread(sock);
			ClientSendThread cst = new ClientSendThread(sock, username);
			
			Thread t1 = new Thread(crt);
			Thread t2 = new Thread(cst);
			
			try {
				//generate pw hash
				MessageDigest md = MessageDigest.getInstance("MD5");
				byte[] digest = md.digest(passBytes);
				String hash = getHexString(digest);
				
				//generate DH exchange params
				AlgorithmParameterGenerator paramgen = AlgorithmParameterGenerator.getInstance("DH");
				paramgen.init(1024); //== 128 bytes
				AlgorithmParameters params = paramgen.generateParameters();
				DHParameterSpec dhParams = (DHParameterSpec)params.getParameterSpec(DHParameterSpec.class);
				
				//generate keys
				KeyPairGenerator keygen = KeyPairGenerator.getInstance("DH");
				keygen.initialize(dhParams, new SecureRandom());
				
				KeyAgreement client_agreement = KeyAgreement.getInstance("DH");
				KeyPair client_keys = keygen.generateKeyPair();
				
				client_agreement.init(client_keys.getPrivate());
				
				//Exchange keys with server
				
				DHMessage dhmsg = new DHMessage("", -1, dhParams.getG(), dhParams.getP(), dhParams.getL(),
						client_keys.getPublic());
				//System.out.println("P: " + dhmsg.getP() + "\nG:" + dhmsg.getG() + "\nL: "+ dhmsg.getL());
				byte[] sermessage = Network.serialize(dhmsg);
				Network.sendBytes(sock, sermessage);
				
				while(sock.getInputStream().available() <=0);
				byte[] respbytes = Network.readBytes(sock);
				
				try{
					Message response = Network.unserialize(respbytes);
					if(response instanceof DHResponseMessage){
						DHResponseMessage dhrm = (DHResponseMessage) response;
						
						//System.out.println("Is DHRM");
						PublicKey server_key = dhrm.getServerPublicKey();		
						client_agreement.doPhase(server_key, true);
						
						SecretKey key = client_agreement.generateSecret("AES");
						//System.out.println("Key is " + getHexString(key.getEncoded()));
						crt.setKey(key);
						cst.setKey(key);
					}
				} catch(ClassNotFoundException e){
					e.printStackTrace();
				}
				
				LoginMessage loginmessage = new LoginMessage(username, -1, hash);
				sermessage = Network.serialize(loginmessage);
				sermessage = Network.encrypt(crt.getKey(), sermessage);
				Network.sendBytes(sock, sermessage);
				
				while(sock.getInputStream().available() <=0);
				
				respbytes = Network.readBytes(sock);
				respbytes = Network.decrypt(crt.getKey(), respbytes);
				
				try {
					Message response = Network.unserialize(respbytes);
					if(response instanceof LoginResponseMessage){
						LoginResponseMessage lrm = (LoginResponseMessage) response;
						if(lrm.getSessionID() == -1L){
							System.out.println("Login failed.  Client exiting.");
							System.exit(0);
						} else {
							System.out.println("Login successful.");
						}
						crt.setSessionID(lrm.getSessionID());
						
						cst.setSessionID(lrm.getSessionID());
					}
				} catch (ClassNotFoundException e) {
					// TODO Auto-generated catch block
					e.printStackTrace();
				}
				
				
			} catch (NoSuchAlgorithmException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
			} catch (Exception e) {
				e.printStackTrace();
			}
			
			t1.start();
			t2.start();
		} catch (NumberFormatException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (UnknownHostException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}

	}
	
	//http://rgagnon.com/javadetails/java-0596.html
	public static String getHexString(byte[] b) {
		  String result = "";
		  for (int i=0; i < b.length; i++) {
		    result +=
		          Integer.toString( ( b[i] & 0xff ) + 0x100, 16).substring( 1 );
		  }
		  return result;
	}

}
